import json
import uuid
from datetime import datetime
from fastapi import APIRouter, HTTPException
from pydantic import BaseModel
from fastapi.responses import StreamingResponse
from typing import Optional, AsyncGenerator

from config import RELEVANCE_THRESHOLD
from embedding import embedding_model
from chroma_client import chroma_client

# 导入新的聊天客户端
from llm_client import OllamaChatClient, generate_chat_stream

from database.operations import (
    get_session,
    create_session,
    update_session,
    add_message,
    get_prompt_template,
    get_session_messages
)

router = APIRouter(tags=["query"], prefix="/api")

# 创建全局聊天客户端实例
chat_client = OllamaChatClient()


class QueryRequest(BaseModel):
    question: str
    session_id: Optional[str] = None
    scenario: str = "knowledge_qa"
    template_id: Optional[str] = None


async def get_prompt_for_scenario(scenario: str, question: str, related_docs: list,
                                  template_id: Optional[str] = None) -> str:
    """获取特定场景的Prompt"""
    # 尝试从数据库获取模板
    template = await get_prompt_template(scenario, template_id)

    if template:
        # 使用数据库中的模板
        prompt = template["content"]
        # 替换变量
        for var in template["variables"]:
            if var == "user_query":
                prompt = prompt.replace(f"{{{var}}}", question)
            elif var == "retrieved_docs":
                prompt = prompt.replace(f"{{{var}}}", "\n".join(related_docs))
        return prompt

    # 如果没有找到模板，使用默认模板
    default_templates = {
        "ai_qa": question,
        "knowledge_qa": f"""已知信息：
{' '.join(related_docs)}

请基于上述信息，回答问题：{question}
要求：1. 只使用提供的信息，不编造内容；2. 语言简洁明了；3. 若信息不足，直接说明。
""",
        "code_review": f"""已知系统代码信息：
{' '.join(related_docs)}

代码改动内容：
{question}

请基于系统现有代码信息，评审上述代码改动：
1. 是否符合现有代码架构和设计模式？
2. 是否存在潜在的兼容性问题或bug？
3. 代码风格是否与现有代码保持一致？
4. 是否有可以优化的地方？
5. 请给出具体的修改建议和理由
""",
        "sql_review": f"""已知数据库表结构信息：
{' '.join(related_docs)}

待评审SQL脚本：
{question}

请基于数据库结构信息，评审上述SQL脚本：
1. SQL语法是否正确？
2. 是否符合现有表结构设计？
3. 是否存在性能问题（如缺少索引、全表扫描等）？
4. 是否有数据一致性或安全性问题？
5. 请给出优化建议和潜在风险提示
""",
        "unit_test": f"""已知代码信息：
{' '.join(related_docs)}

代码改动内容：
{question}

请基于上述代码，生成单元测试代码：
1. 使用Python的unittest或pytest测试框架
2. 覆盖主要功能点和边界条件
3. 包含正常场景和异常场景测试
4. 测试代码应包含清晰的注释
5. 给出测试用例设计思路
"""
    }

    return default_templates.get(scenario, default_templates["knowledge_qa"])


async def build_conversation_messages(session_id: str, current_question: str,
                                      related_docs: list, scenario: str,
                                      template_id: Optional[str] = None) -> list:
    """
    构建包含对话历史和当前问题的消息列表

    Args:
        session_id: 会话ID
        current_question: 当前问题
        related_docs: 相关文档
        scenario: 场景类型
        template_id: 模板ID

    Returns:
        消息列表，格式为 [{"role": "user/assistant/system", "content": "..."}]
    """
    # 获取历史消息
    history_messages = await get_session_messages(session_id)

    # 构建系统提示词
    system_prompt = await build_system_prompt(scenario, current_question, related_docs, template_id)

    # 构建消息列表
    messages = []

    # 添加系统提示词（如果有）
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})

    # 添加历史对话（排除当前刚添加的用户消息）
    for msg in history_messages:
        # 跳过当前刚添加的用户消息
        if msg.get("content") == current_question and msg.get("is_user"):
            continue

        role = "user" if msg.get("is_user") else "assistant"
        messages.append({"role": role, "content": msg.get("content", "")})

    # 添加当前用户消息
    messages.append({"role": "user", "content": current_question})

    return messages


async def build_system_prompt(scenario: str, question: str, related_docs: list,
                              template_id: Optional[str] = None) -> str:
    """
    构建系统提示词

    Args:
        scenario: 场景类型
        question: 问题
        related_docs: 相关文档
        template_id: 模板ID

    Returns:
        系统提示词
    """
    # 如果有相关文档，构建包含文档的系统提示
    if related_docs:
        base_prompt = await get_prompt_for_scenario(scenario, question, related_docs, template_id)

        # 根据场景调整系统提示词
        if scenario == "knowledge_qa":
            return f"""你是一个有帮助的AI助手。请基于以下已知信息回答问题：

{base_prompt}

请保持对话的连贯性，如果用户的问题与之前的对话相关，请结合上下文进行回答。"""

        elif scenario == "code_review":
            return f"""你是一个专业的代码评审助手。请基于以下代码信息进行评审：

{base_prompt}

请给出具体、可操作的改进建议。"""

        elif scenario == "sql_review":
            return f"""你是一个专业的SQL评审助手。请基于以下数据库信息进行SQL评审：

{base_prompt}

请指出潜在的性能问题和安全隐患。"""

        elif scenario == "unit_test":
            return f"""你是一个专业的单元测试生成助手。请基于以下代码信息生成测试：

{base_prompt}

请确保测试覆盖主要功能和边界情况。"""

    # 如果没有相关文档，使用通用系统提示
    if scenario == "ai_qa":
        return "你是一个有帮助的AI助手。请直接回答用户的问题，保持对话的连贯性。"
    else:
        return f"你是一个有帮助的AI助手。请回答用户关于{scenario}的问题，保持对话的连贯性。"


@router.post("/query/stream")
async def query_knowledge_stream(request: QueryRequest):
    """处理用户查询（流式响应版本）"""

    # 1. 使用带分数的重排序方法检索相关文档
    scored_results = chroma_client.query_with_scores(
        query_text=request.question,
        n_results=3
    )

    # 2. 判断是否有相关知识
    has_relevant_knowledge = False
    related_docs = []

    if scored_results:
        # 检查最高分是否超过阈值
        top_score = scored_results[0]["reranker_score"]
        if top_score >= RELEVANCE_THRESHOLD:
            has_relevant_knowledge = True
            related_docs = [result["content"] for result in scored_results
                            if result["reranker_score"] >= RELEVANCE_THRESHOLD]

    # 3. 处理会话
    session_id = request.session_id or str(uuid.uuid4())
    current_time = datetime.now().isoformat()

    # 检查会话是否存在
    session = await get_session(session_id)
    if not session:
        # 创建新会话
        await create_session(session_id, request.question)
    else:
        # 更新现有会话
        await update_session(session_id, request.question)

    # 保存用户消息
    await add_message(session_id, request.question, is_user=True)

    # 4. 构建消息列表（使用聊天模式）
    if has_relevant_knowledge:
        # 有相关知识：使用完整场景
        messages = await build_conversation_messages(
            session_id=session_id,
            current_question=request.question,
            related_docs=related_docs,
            scenario=request.scenario,
            template_id=request.template_id
        )
    else:
        # 没有相关知识：使用ai_qa场景
        messages = await build_conversation_messages(
            session_id=session_id,
            current_question=request.question,
            related_docs=[],
            scenario='ai_qa',
            template_id=request.template_id
        )

    # 5. 创建流式响应
    async def generate() -> AsyncGenerator[str, None]:
        full_response = ""
        # 调用聊天模式的流式LLM
        async for chunk in generate_chat_stream(messages):
            full_response += chunk
            # 发送当前块
            yield f"data: {json.dumps({'chunk': chunk, 'done': False})}\n\n"

        # 流式响应完成后保存AI回复
        await add_message(session_id, full_response, is_user=False)

        # 发送最终完成信号
        yield f"data: {json.dumps({
            'done': True,
            'session_id': session_id,
            'related_documents': related_docs
        })}\n\n"

    return StreamingResponse(
        generate(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "Access-Control-Allow-Origin": "*",
        }
    )


# 会话管理接口
@router.get("/sessions/{session_id}/messages")
async def get_session_messages_api(session_id: str):
    """获取指定会话的所有消息（API接口）"""
    try:
        messages = await get_session_messages(session_id)
        return {
            "session_id": session_id,
            "messages": [
                {
                    "id": msg.get("id"),
                    "content": msg.get("content"),
                    "is_user": msg.get("is_user"),
                    "created_at": msg.get("created_at")
                }
                for msg in messages
            ]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取消息失败: {str(e)}")